# -*- coding: utf-8 -*-
"""
 * @file main.py
 * @author SDK_TEAM
 * @brief
 * @version 0.1
 * @date 2022-11-3
 * Copyright:
 * © 2018 北京灵汐科技有限公司 版权所有。
 * 注意：以下内容均为北京灵汐科技有限公司原创，未经本公司允许，不得转载，否则将视为侵权；对于不遵守此声明或者其他违法使用以下内容者，本公司依法保留追究权。
 * © 2018 Lynxi Technologies Co., Ltd. All rights reserved.
 * NOTICE: All information contained here is, and remains the property of Lynxi. This file can not
 * be copied or distributed without the permission of Lynxi Technologies Co., Ltd.
"""

import argparse
import pylynchipsdk as sdk
import numpy as np
import threading
import _thread
import struct
from ctypes import *
from typing import cast
import getopt
import ipeParamYolov5s6

import os
import sys
import time
sys.path.append('../../../')
import common.python.common as common
import common.python.bufferpool as bufferpool
from common.python.plugin_utils import *
from common.python.common import *
from common.python import image_decoder
from common.python import image_encoder

def main():
    parser = argparse.ArgumentParser(description="yolov5s6")

    parser.add_argument(
        "input_image",
        metavar="input_image",
        type=str,
        default="../../../data/yolov5s6_detect_with_image_sample/test.jpeg",
        help="input jpeg file path",
    )
    parser.add_argument(
        "output_image",
        metavar="output_image",
        type=str,
        default="./output.jpeg",
        help="output jpeg file path",
    )
    parser.add_argument(
        "device_id", metavar="device_id", type=int, default=0, help="lynxi device id"
    )
    parser.add_argument(
        "-m",
        "--model",
        metavar="model",
        type=str,
        default="../../../model/yolov5s6_b1_h1280_w1280_c3_iuint8_ofloat16_0000_x7_quant/Net_0",
        help="yolov5s6 model path",
    )
    parser.add_argument(
        "-p",
        "--post_plugin",
        metavar="post_plugin",
        type=str,
        default="../../../common/plugin/postprocess/cross_compiled_lib/libYolov5n6X7PostPlugin.so",
        help="post process plugin path",
    )
    parser.add_argument(
        "-o",
        "--osd_plugin",
        metavar="osd_plugin",
        type=str,
        default="../../../common/plugin/osd/lib/libOsdPlugin.so",
        help="osd plugin path",
    )
    parser.add_argument(
        "-l",
        "--label_file",
        default="../../../common/plugin/postprocess/yolox/label/label.v5x7.json",
        type=str,
        help="label file",
    )

    args = parser.parse_args()

    input_file = args.input_image
    output_file = args.output_image
    model_path = args.model
    device_id = args.device_id
    post_plugin_path = args.post_plugin
    osd_plugin_path = args.osd_plugin

    if not os.path.exists(input_file):
        raise ("error: input jpeg file does not exist.")

    if not is_valid_device(device_id):
        raise ("error: device id is invalid.")

    if not os.path.exists(model_path):
        raise ("error: model does not exist.")

    if not os.path.exists(post_plugin_path):
        raise ("error: post process plugin does not exist.")

    if not os.path.exists(osd_plugin_path):
        raise ("error: osd plugin does not exist.")

    # 创建 context
    context, ret = sdk.lyn_create_context(device_id)
    common.error_check(ret, "lyn_create_context")
    ret = sdk.lyn_register_error_handler(common.default_stream_error_handle, None)
    common.error_check(ret, "lyn_register_error_handler")
    # 创建 stream
    stream, ret = sdk.lyn_create_stream()
    common.error_check(ret, "lyn_create_stream")

    # 初始化解码图片类
    print("image decode")

    image_dec = image_decoder.ImageDecoder()
    image_dec.init(input_file, False)
    img_info = image_dec.get_image_dec_info()

    # 同步解码图片到 Device 侧
    decode_img = sdk.lyn_codec_buf_t()
    decode_img.size = img_info.output.predict_buf_size
    p_decode_buf_out, ret = sdk.lyn_malloc(decode_img.size)
    error_check(ret)
    decode_img.data = p_decode_buf_out
    decodeImg_data_ptr = pythonapi.PyCapsule_GetPointer(decode_img.data, None)
    image_dec.decode_image_to_device(stream, decode_img.data)

    # 加载模型
    model = ModelInfo(model_path)

    # 加载 Plugin
    osd_plugin, ret = sdk.lyn_plugin_register(osd_plugin_path)
    common.error_check(ret, "lyn_plugin_register")
    post_plugin, ret = sdk.lyn_plugin_register(post_plugin_path)
    common.error_check(ret, "lyn_plugin_register")

    label_list, label_info_dev = create_label_list_with_names(
        args.label_file
    )

    print("ipe and infer")
    pDetectIpeBuf, ret = sdk.lyn_malloc(model.input_size)
    error_check(ret)
    sdk.lyn_memset(pDetectIpeBuf, 0, model.input_size)
    ipe = ipeParamYolov5s6.IpeParamYolov5s6(
        model.width, model.height
    )
    ipe.set_img_info(
        img_info.output.width,
        img_info.output.height,
        sdk.lyn_pixel_format_t.LYN_PIX_FMT_NV12,
    )
    ipe.calc_param(stream, decode_img.data, pDetectIpeBuf)

    # 调用 APU 推理接口
    pDevDetectApuBuf, ret = sdk.lyn_malloc(model.output_size)
    error_check(ret)
    error_check(
        sdk.lyn_execute_model_async(
            stream,
            model.model,
            pDetectIpeBuf,
            pDevDetectApuBuf,
            model.batch_size,
        )
    )

    pythonapi.PyCapsule_GetPointer.restype = c_void_p
    pythonapi.PyCapsule_GetPointer.argtypes = [py_object, c_char_p]

    pDevDetectApuBuf_ptr = pythonapi.PyCapsule_GetPointer(pDevDetectApuBuf, None)

    boxes_info_dev, ret = sdk.lyn_malloc(ctypes.sizeof(Box))
    common.error_check(ret)
    boxes_info_dev_ptr = pythonapi.PyCapsule_GetPointer(boxes_info_dev, None)

    print("post process")

    post_para = YoloxPostProcessInfo_t()
    post_para.height = model.height
    post_para.width = model.width
    post_para.ori_height = img_info.output.height
    post_para.ori_width = img_info.output.width
    post_para.score_threshold = 0.25
    post_para.nms_threshold = 0.45
    post_para.nms_top_k = 500
    post_para.is_pad_resize = 1
    post_para.output_tensor = pDevDetectApuBuf_ptr
    post_para.class_num = label_list.labelNum
    post_para.labelList = label_list
    post_para.boxesInfo = boxes_info_dev_ptr
    post_para.anchorSize = 3

    yolo_data = ctypes.string_at(
        ctypes.addressof(post_para), ctypes.sizeof(post_para)
    )
    ret = sdk.lyn_plugin_run_async(
        stream,
        post_plugin,
        "lynYoloxPostProcess",
        yolo_data,
        ctypes.sizeof(YoloxPostProcessInfo_t),
    )
    common.error_check(ret, "lyn_plugin_run_async")

    print("osd")

    osd_para = struct.pack(
        "1P3i1P4i",
        boxes_info_dev_ptr,
        img_info.output.width,
        img_info.output.height,
        sdk.lyn_pixel_format_t.LYN_PIX_FMT_NV12,
        decodeImg_data_ptr,
        24,
        4,
        2,
        4,
    )
    common.error_check(
        sdk.lyn_plugin_run_async(
            stream,
            osd_plugin,
            "lynDrawBoxAndText",
            osd_para,
            len(osd_para),
        )
    )

    error_check(sdk.lyn_synchronize_stream(stream))

    print("save new image")
    encoder = image_encoder.ImageEncoder()
    encoder.encode_image_and_save(stream, output_file, decode_img, img_info)

    ipe.destory()

    # 等待流处理完
    ret = sdk.lyn_synchronize_stream(stream)
    common.error_check(ret, "lyn_synchronize_stream")
    ret = sdk.lyn_destroy_stream(stream)
    common.error_check(ret, "lyn_destroy_stream")


    sdk.lyn_free(label_info_dev)
    sdk.lyn_free(pDetectIpeBuf)
    sdk.lyn_free(pDevDetectApuBuf)
    sdk.lyn_free(decode_img.data)
    sdk.lyn_free(boxes_info_dev)

    # 卸载模型
    ret = sdk.lyn_plugin_unregister(post_plugin)
    common.error_check(ret, "lyn_plugin_unregister")
    ret = sdk.lyn_plugin_unregister(osd_plugin)
    common.error_check(ret, "lyn_plugin_unregister")
    model.unload_model()
    # 销毁 context
    ret = sdk.lyn_destroy_context(context)
    common.error_check(ret, "lyn_destroy_context")


if __name__ == "__main__":
    main()